import numpy as np
import torch
def one_hot_encode(label,n_class,cuda):
    # print(n_class)
    # label = label[np.newaxis, :, :]  # add new dim in any dim
    # print(label.shape)

    # (1, 530, 730) -> [1, 530, 730, 1]
    label = label.unsqueeze(3)
    label = label.long()
    # print(label.shape)

    # [1, 530, 730, 1] -> [1, 530, 730, 38]
    zeros=torch.zeros(label.shape[0], label.shape[1], label.shape[2], n_class).long()
    if cuda:
        zeros= zeros.cuda()
    label = zeros.scatter_(3, label, 1).long()
    return label


if __name__ == '__main__':
    a = np.random.randn(1,2,3)
    print(a)

